from graphviz import Digraph

def wrap_text(text, max_length):
    """根据最大长度对文本进行换行"""
    words = text.split()
    lines = []
    current_line = ""
    
    for word in words:
        if len(current_line) + len(word) + 1 <= max_length:
            current_line += (" " + word if current_line else word)
        else:
            lines.append(current_line)
            current_line = word
    if current_line:
        lines.append(current_line)
    
    return "\n".join(lines)

def create_mindmap(title, sections, subsections, path, max_length=30, wrap_levels=None):
    if wrap_levels is None:
        wrap_levels = []  # 默认只有第一级换行

    dot = Digraph(format='png')
    dot.attr(dpi='300')
    dot.attr(rankdir='LR', splines='spline')  # 改回流畅曲线
    dot.attr('node', shape='box', style='rounded, filled', fontname='Arial', penwidth='2')
    
    # 颜色设置（低饱和度，优雅风格）
    color_primary = '#EDEFF0'   # 一级节点（低饱和度红）
    border_primary = '#B07C7C'  # 一级节点边框（深一点的红）
    
    color_secondary = '#F2F6F9' # 二级节点（低饱和度蓝）
    border_secondary = '#7C95B0' # 二级节点边框（深一点的蓝）
    
    color_tertiary = '#FFF1F7' 
    border_tertiary = '#FEDBEB' 

    # 一级节点
    title_wrapped = wrap_text(title, max_length) if 0 in wrap_levels else title
    dot.node('0', title_wrapped, fillcolor=color_primary, color=border_primary)
    
    # 二级节点
    for key, label in sections.items():
        label_wrapped = wrap_text(label, max_length) if 1 in wrap_levels else label
        dot.node(key, label_wrapped, fillcolor=color_secondary, color=border_secondary)
        dot.edge('0', key, tailport='e', headport='w', arrowhead='none')
    
    # 三级节点
    for key, (parent, label) in subsections.items():
        label_wrapped = wrap_text(label, max_length) if 2 in wrap_levels else label
        dot.node(key, label_wrapped, fillcolor=color_tertiary, color=border_tertiary)
        dot.edge(parent, key, tailport='e', headport='w', arrowhead='none')
    
    # 生成并保存图像
    dot.render(path, view=False)


if __name__ == "__main__":
    # 二级节点
    sections = {
        '1': '1. Introduction',
        '2': '2. Background and Definitions',
        '3': '3. P vs NP and NP-Completeness',
        '4': '4. Polynomial Hierarchy',
        '5': '5. Space Complexity',
        '6': '6. Oracle Separation and Relativization',
        '7': '7. Time Complexity and Circuit Complexity',
        '8': '8. Conclusion'
    }
    # 三级节点
    subsections = {
        '3.1': ('3', '3.1 Historical Context and Significance of P vs NP'),
        '3.2': ('3', '3.2 NP-Completeness and the Cook-Levin Theorem'),
        '3.3': ('3', '3.3 Applications and Implications of NP-Completeness'),
        '3.4': ('3', '3.4 Current Status and Open Questions'),
        '5.1': ('5', '5.1 Foundational Concepts and Measurement'),
        '5.2': ('5', '5.2 Space Complexity in Proof Systems'),
        '5.3': ('5', '5.3 Space Complexity in Quantum and Classical Algorithms'),
        '5.4': ('5', '5.4 Applications and Implications in Algorithm Design'),
        '6.1': ('6', '6.1 Foundations of Oracle Separation'),
        '6.2': ('6', '6.2 Relativization and Its Limitations'),
        '6.3': ('6', '6.3 Advanced Techniques in Oracle Separation'),
        '6.4': ('6', '6.4 Applications in Quantum Complexity')
    }
    create_mindmap(
        "Title",
        sections,
        subsections,
        path="complexity_mindmap",
        max_length=20,  # 设置每行最大长度
        wrap_levels=[0, 1, 2]  # 控制哪些级别换行
    )